'''这里
'''

import operator as op
import numpy as np
from misc import *


class _BaseTree:
    '''作为KDTree和BallTree的父类，不直接使用'''

    def __init__(self, leafsize, distance):
        (self.leafsize, self.distance) = (leafsize, distance)

    def create(self, xs):
        '''根据输入的点集，构造树形结构'''
        self.xs = xs.copy()
        self.tree = self._create(range(len(xs)))

    def find_nearests(self, x, k):
        '''找出与x最接近的k个点'''

        nearests = [[np.inf, -1] for i in range(k)]
        self._find_nearests(x, self.tree, nearests)

        return nearests

    def _create(self, rows):
        if len(rows) > self.leafsize:
            (crit, l_rows, r_rows) = self._split_rows(rows)
            return [crit, self._create(l_rows), self._create(r_rows)]
        else:
            return rows

    def _split_rows(self, rows):
        pass

    def _find_nearests(self, x, tree, nearests):
        if self._is_leaf(tree):
            self._find_nearest_on_leaf(x, tree, nearests)
        else:
            self._find_nearests_on_branch(x, tree, nearests)

    def _find_nearests_on_branch(self, x, tree, nearests):
        assert(self._is_branch(tree))

        if self._is_near_to_left_tree(x, tree[0]):
            self._find_nearests(x, tree[1], nearests)
            if not self._can_rule_out_right_tree(x, tree[0], nearests):
                self._find_nearests(x, tree[2], nearests)

        else:
            self._find_nearests(x, tree[2], nearests)
            if not self._can_rule_out_left_tree(x, tree, nearests):
                self._find_nearests(x, tree[1], nearests)

    def _is_near_to_left_tree(self, x, cert):
        pass

    def _can_rule_out_left_tree(self, x, cert, nearests):
        pass

    def _find_nearest_on_leaf(self, x, tree, nearests):
        assert(self._is_leaf(tree))

        ds = self.distance(x, self.xs[tree, :])
        for i in range(len(tree)):
            pos = len(nearests)
            while pos > 0:
                if ds[i] < nearests[pos-1][0]:
                    pos -= 1
                else:
                    break

            if pos < len(nearests):
                nearests.pop()
                nearests.insert(pos, [ds[i], tree[i]])

    def _is_leaf(self, tree):
        return len(tree) != 3 or type(tree[2]) != list

    def _is_branch(self, tree):
        return not self._is_leaf(tree)

class KDTree(_BaseTree):
    def __init__(self, leafsize):
        super().__init__(leafsize, euclidean)
    
    def _split_rows(self, rows):
        
        var = np.var(self.xs[rows, :], 0)
        col = var.argsort()[-1]

        p = self.xs[rows, col].argsort()
        s = np.ceil(len(p) / 2)

        return [[col, self.xs[rows[p[s]],:]], [rows[i] for i in p[:s]], [rows[i] for i in p[s:]]]

    def _is_near_to_left_tree(self, x, cert):
        return x[cert[0]] <= cert[1]

    def _can_rule_out_left_tree(self, x, tree, nearests):
        return nearests[-1][0] > x[cert[0]] - cert[1]

    def _can_rule_out_right_tree(self, x, cert, nearests):
        return nearests[-1][0] < cert[1] - x[cert[0]]

class BallTree(_BaseTree):
    def __init__(self, leafsize, distance=euclidean):
        super().__init__(leafsize, euclidean)

    def _split_rows(self, rows):
        ds = self.distance(self.xs[rows[0]], self.xs[rows])

        si = ds.argsort()

        m = np.ceil(len(si)/2)

        cert = [rows[0], ds[si[m]]]
        l_rows = [rows[i] for i in si[:m]]
        r_rows = [rows[i] for i in si[m:]]

        return (cert, l_rows, r_rows)

    def _is_near_to_left_tree(self, x, cert):
        return self.distance(x, self.xs[cert[0],:]) <= cert[1]

    def _can_rule_out_left_tree(self, x, cert, nearests):
        return self.distance(x, self.xs[cert[0],:]) - nearests[-1][0] > cert[1]

    def _can_rule_out_right_tree(self, x, cert, nearests):
        return nearests[-1][0] + self.distance(x, self.xs[cert[0],:]) <= cert[1]





def _test_kd_ball_tree():
    xs = np.array([1,2,3,4,5,6,7,8,9],ndmin=2).transpose()
    print(xs)


    #tree = BallTree(5)
    tree = KDTree(5)

    tree.create(xs)
    nearests = tree.find_nearests(np.array([4.8],ndmin=2), 2)
    print(tree.tree)
    print(nearests)



if __name__ == "__main__":
    _test_kd_ball_tree()


